import pygp
import math

NUMFITCASE = 30

class SymReg:
    '''
    Symbolic regression
    '''
    def __init__(self, rand):
        self.random = rand
        self.fitcase = []
        self.SuccPredicate = False
        self.term = pygp.Term(rand)
        self.LoadTerm()
        self.SetFitCase()

    def LoadTerm(self):
        print 'Loading terms...'
        self.term.AddTerm("ADD", 2)
        self.term.AddTerm("SUB", 2)
        self.term.AddTerm("DIV", 2)
        self.term.AddTerm("MUL", 2)
        self.term.AddTerm("SIN", 1)
        self.term.AddTerm("COS", 1)
        self.term.AddTerm("EP", 1)
        self.term.AddTerm("X", 0)
        self.term.Swap(100)
        print 'Done!'

    def SetFitCase(self):
        for i in range(NUMFITCASE):
            fitcase = [2*self.random.random() - 1, 0]
            x = fitcase[0]
            #function
            y = x**4 + x**3 + x**2 + x
            #y = math.sin(x) + math.cos(x)
            #y = x**3 + x**2 + x
            fitcase[1] = y
            self.fitcase.append(fitcase)

    def Compute(self, node, x):
        '''
        compute the function (in tree form) with given input
        @node: tree root
        @x: value of x
        '''
        if node.name == 'X':
            return x
        else:
            if node.name == 'ADD':
                l = self.Compute(node.children[0], x)
                r = self.Compute(node.children[1], x)
                return l + r
            elif node.name == 'SUB':
                l = self.Compute(node.children[0], x)
                r = self.Compute(node.children[1], x)
                return l - r
            elif node.name == 'MUL':
                l = self.Compute(node.children[0], x)
                r = self.Compute(node.children[1], x)
                return l * r
            elif node.name == 'DIV':
                l = self.Compute(node.children[0], x)
                r = self.Compute(node.children[1], x)
                if r == 0:
                    return 1
                else:
                    return l / r
            elif node.name == 'SIN':
                l = self.Compute(node.children[0], x)
                try:
                    return math.sin(l)
                except ValueError:
                    return math.exp(700)
            elif node.name == 'COS':
                l = self.Compute(node.children[0], x)
                try:
                    return math.cos(l)
                except ValueError:
                    return math.exp(700)
            elif node.name == 'EP':
                l = self.Compute(node.children[0], x)
                try:
                    return math.exp(l)
                except OverflowError:
                    if l > 0:
                        return math.exp(700)
                    else:
                        return -math.exp(700)
            elif node.name == 'LOG':
                l = self.Compute(node.children[0], x)
                if l == 0:
                    return 0
                else:
                    return math.log10(math.fabs(l))

    def ComputeRF(self, individual):
        '''
        Compute the row fitness for an individual
        '''
        #print 'computeRF...'
        summ = 0
        hit  = 0
        for i in range (NUMFITCASE):
            t = math.fabs(self.Compute(individual.chrom, self.fitcase[i][0])
                          -self.fitcase[i][1])
            if t <= 0.01:
                hit += 1
            summ += t
        if (hit == NUMFITCASE):
            self.SuccPredicate = True
        return summ
 
 
